from transmission_simulator.structure.internal_model import InternalModel
from transmission_simulator.io.info_loader import InfoLoader
from transmission_simulator.io.results_writer import ResultsWriter


class Controller:
    def __init__(self, model_func=None, input_path="./input"):
        # 负责读取数据文件
        self._loader = InfoLoader(input_path)
        # 负责处理结果
        self._writer = ResultsWriter()
        # 负责地区内部传染模拟
        self._model = InternalModel(model_func)
        # 当前日期
        self._cur_day = 0
        # 地区数据初始化
        self._regions = self._loader.get_init_info()
        # 将初始信息写入结果buffer
        self._write_cur_day()

    def _hard_update(self):
        cur_update = self._loader.get_cur_update(self._cur_day)
        for info in cur_update:
            # info['name']代表地名，_regions[info['name']]代表地名对应的Region对象
            self._regions[info['name']].update(info)

    def _migration_collect(self):
        migration_list = []
        # 遍历所有Region类
        for region in self._regions.values():
            migration_list.append(region.get_migration())
        return migration_list

    def _model_update(self):
        # 遍历所有Region类
        for region in self._regions.values():
            # 从region获取数据
            population_dict = region.get_population()
            model_param_dict = region.get_param()
            # 计算次日人口组成
            update_info = self._model.calculate(population_dict, model_param_dict)
            # 更新
            region.update(update_info)

    def _migration_update(self, migration_list):
        for region_migration in migration_list:
            for dst_region in region_migration:
                # dst_region['name']代表输出目标地名
                self._regions[dst_region['name']].update(dst_region)

    def _write_cur_day(self):
        results_dict = {}
        # 遍历所有Region类
        for region_item in self._regions.items():
            # 保存region对象的数据
            detail = {'name': region_item[0]}
            detail.update(region_item[1].get_population())
            detail.update(region_item[1].get_export_info())
            detail.update(region_item[1].get_param())
            # 地名为key，地区数据为value
            detail_dict = {region_item[0]: detail}
            results_dict.update(detail_dict)
        # 写入buffer
        self._writer.write_to_buffer(results_dict, self._cur_day)

    def next_day(self):
        # 日期移动
        self._cur_day += 1
        # 使用外部数据更新
        self._hard_update()
        # 计算地区内部传染情况并更新
        self._model_update()
        # 计算迁出人口组成
        migration_list = self._migration_collect()
        # 更新迁出人口组成
        self._migration_update(migration_list)
        # 将结果写入缓存
        self._write_cur_day()
        return

    def set_model_func(self, model_func):
        # 更新自定义函数
        self._model.set_func(model_func)

    def write(self, path="./output"):
        # 写入到文件
        self._writer.write_to_file(path)

    def reimport_update(self):
        return self._loader.reimport_update_json()

    def get_results(self, region_name=None, day=None):
        # 获得buffer
        buffer = self._writer.get_buffer()
        # 地名为空
        if region_name is None:
            # 地名、天数都为空
            if day is None:
                return buffer
            # 地名为空，天数非空
            else:
                return buffer.get(str(day))
        # 地名非空
        else:
            # 地名非空，天数为空
            if day is None:
                # 检查地名key值是否存在
                if region_name not in self._regions.keys():
                    return
                results = {}
                # 遍历每一天
                for day_item in buffer.items():
                    # 天数为key，地区数据为value
                    detail_dict = {day_item[0]: day_item[1][region_name]}
                    results.update(detail_dict)
                return results
            # 地名、天数都非空
            else:
                day_result = buffer[str(day)]
                # 遍历各个地区结果
                return day_result[region_name]
